from typing import TypeVar, Generic, Optional, List, Dict, Any, AsyncGenerator
from pydantic import BaseModel, Field, ConfigDict
from contextlib import asynccontextmanager
from inspect import isawaitable
from abc import ABC, abstractmethod
from datetime import datetime, timedelta

import asyncio
import logging
import traceback

# 使用泛型类型变量
E = TypeVar("E")


class PoolConfig:
    """连接池配置类，提供更灵活的配置管理"""

    def __init__(
        self,
        max_connections: int = 10,
        min_connections: int = 1,
        connection_timeout: float = 30.0,
        idle_timeout: float = 300.0,
        connection_retry: int = 3,
        retry_interval: float = 1.0,
        check_interval: int = 120,
        logger: Optional[logging.Logger] = None,
    ):
        """
        初始化连接池配置

        Args:
            max_connections: 连接池中允许的最大连接数
            min_connections: 连接池中保持的最小连接数
            connection_timeout: 连接超时时间（秒）
            idle_timeout: 空闲连接的超时时间（秒）
            connection_retry: 连接失败时的重试次数
            retry_interval: 重试间隔时间（秒）
            check_interval: 后台维护任务的检查间隔时间（秒）
            logger: 自定义日志记录器
        """
        self.max_connections = max_connections
        self.min_connections = min_connections
        self.connection_timeout = connection_timeout
        self.idle_timeout = idle_timeout
        self.connection_retry = connection_retry
        self.retry_interval = retry_interval
        self.check_interval = check_interval
        self.logger = logger or logging.getLogger(__name__)


class Entity(BaseModel, Generic[E]):
    """
    连接实体类，管理单个连接的元数据
    使用泛型以支持不同类型的连接客户端
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    client: E = Field(description="客户端连接对象")
    created_at: datetime = Field(
        default_factory=datetime.now, description="实体创建时间"
    )
    last_used: datetime = Field(
        default_factory=datetime.now, description="实体最后使用时间"
    )
    in_use: bool = Field(default=False, description="实体是否正在使用")
    health_check_failures: int = Field(default=0, description="健康检查失败次数")


class BasePool(ABC, Generic[E]):
    """
    抽象连接池基类，提供通用的连接池管理功能
    支持异步操作和灵活配置
    """

    def __init__(self, config: Optional[PoolConfig] = None):
        """
        初始化连接池

        Args:
            config: 连接池配置对象，如果未提供则使用默认配置
        """
        # 使用默认配置或传入的配置
        self.config = config or PoolConfig()

        # 连接池管理属性
        self._pool: List[Entity[E]] = []  # 存储连接对象及其元数据
        self._lock = asyncio.Lock()  # 用于同步连接池操作的锁
        self._initialization_lock = asyncio.Lock()  # 用于初始化的锁

        # 池状态管理
        self._initialized = False  # 连接池是否已初始化
        self._closing = False  # 连接池是否正在关闭
        self._maintenance_task = None  # 后台维护任务

    @abstractmethod
    async def create_connection(self) -> Optional[Entity[E]]:
        """
        抽象方法：创建新的连接
        子类必须实现具体的连接创建逻辑

        Returns:
            创建的连接实体，如果创建失败则返回None
        """
        raise NotImplementedError("子类必须实现 create_connection 方法")

    @abstractmethod
    async def validate_connection(self, connection: Entity[E]) -> bool:
        """
        抽象方法：验证连接的有效性
        子类必须实现具体的连接验证逻辑

        Args:
            connection: 待验证的连接实体

        Returns:
            连接是否有效
        """
        raise NotImplementedError("子类必须实现 validate_connection 方法")

    @asynccontextmanager
    async def acquire(self, timeout: Optional[float] = None) -> AsyncGenerator[E, None]:
        """
        获取连接的上下文管理器，支持超时和重试机制

        Args:
            timeout: 获取连接的最大等待时间，默认使用配置中的超时时间

        Returns:
            可用的连接客户端

        Raises:
            TimeoutError: 无法获取可用连接
        """
        # 如果未初始化，先初始化连接池
        if not self._initialized:
            await self.initialize()

        # 使用配置的超时时间或传入的超时时间
        max_wait_time = timeout or self.config.connection_timeout
        wait_time = 0
        retry_interval = self.config.retry_interval

        connection = None
        while wait_time < max_wait_time:
            try:
                connection = await self._get_connection()
                if connection:
                    break
            except Exception as e:
                self.config.logger.warning(f"获取连接时发生错误: {e}")

            self.config.logger.info(f"等待可用连接... ({wait_time}/{max_wait_time}秒)")
            await asyncio.sleep(retry_interval)
            wait_time += retry_interval

        if not connection:
            raise TimeoutError("无法获取可用连接，请稍后重试")

        try:
            # 提供连接给调用者
            yield connection.client
        finally:
            # 使用完毕后释放连接
            await self._release_connection(connection)

    async def initialize(self):
        """
        初始化连接池，创建最小数量的连接
        使用锁确保线程安全的初始化
        """
        async with self._initialization_lock:
            if self._initialized:
                return

            # 创建最小连接数量的连接
            for _ in range(self.config.min_connections):
                try:
                    new_conn = await self.create_connection()
                    if new_conn:
                        self._pool.append(new_conn)
                except Exception as e:
                    self.config.logger.error(f"初始化连接失败: {e}")

            # 启动后台维护任务
            self._maintenance_task = asyncio.create_task(self._maintain_pool())
            self._initialized = True
            self.config.logger.info(
                f"{str(type(self._pool[0].client)).split('.')[-1]}连接池初始化完成"
            )

    async def _get_connection(self) -> Optional[Entity[E]]:
        """
        从连接池获取一个可用且有效的连接

        Returns:
            可用的连接对象，如果没有可用连接则返回None
        """
        async with self._lock:
            # 首先查找空闲且有效的连接
            for conn in self._pool:
                if not conn.in_use and await self.validate_connection(conn):
                    conn.in_use = True
                    conn.last_used = datetime.now()
                    return conn

            # 如果没有空闲连接并且未达到最大连接数，则创建新连接
            if len(self._pool) < self.config.max_connections:
                try:
                    new_conn = await self.create_connection()
                    if new_conn:
                        new_conn.in_use = True
                        self._pool.append(new_conn)
                        return new_conn
                except Exception as e:
                    self.config.logger.error(f"创建新连接失败: {e}")

            # 如果已达到最大连接数且没有空闲连接，返回None
            self.config.logger.warning("无可用连接，已达到最大连接数")
            return None

    async def _release_connection(self, connection: Entity[E]):
        """
        释放连接回连接池

        Args:
            connection: 要释放的连接实体
        """
        if self._closing:
            # 如果连接池正在关闭，关闭连接而不是释放
            await self._close_connection(connection)
            return

        async with self._lock:
            if connection in self._pool:
                connection.in_use = False
                connection.last_used = datetime.now()
                self.config.logger.debug("连接已释放回连接池")

    async def _close_connection(self, connection: Entity[E]):
        """
        关闭并移除一个连接

        Args:
            connection: 要关闭的连接实体
        """
        try:
            # 具体的关闭逻辑由子类实现
            client = connection.client
            if hasattr(client, "close"):
                if isawaitable(client.close):
                    await client.close()
                else:
                    client.close()
            self.config.logger.debug("连接已关闭")
        except Exception as e:
            self.config.logger.error(f"关闭连接出错: {traceback.format_exc()}")
        finally:
            async with self._lock:
                if connection in self._pool:
                    self._pool.remove(connection)

    async def _maintain_pool(self):
        """
        后台任务：维护连接池
        - 关闭超时的空闲连接
        - 检查并修复不健康的连接
        - 确保连接池中至少有最小数量的连接
        """
        while not self._closing:
            try:
                await asyncio.sleep(self.config.check_interval)

                # 清理空闲超时连接
                await self._cleanup_idle_connections()

                # 修复不健康的连接
                await self._repair_unhealthy_connections()

                # 补充最小连接数
                await self._ensure_min_connections()

            except Exception as e:
                self.config.logger.error(
                    f"连接池维护任务出错: {traceback.format_exc()}"
                )

    async def _cleanup_idle_connections(self):
        """清理空闲超时的连接"""
        async with self._lock:
            current_time = datetime.now()
            idle_connections = [
                conn
                for conn in self._pool
                if not conn.in_use
                and (current_time - conn.last_used)
                > timedelta(seconds=self.config.idle_timeout)
            ]

            # 计算可以关闭的连接数量（保持最小连接数）
            total_count = len(self._pool)
            idle_count = len(idle_connections)
            close_count = min(idle_count, total_count - self.config.min_connections)

            if close_count > 0:
                self.config.logger.info(f"关闭 {close_count} 个空闲超时连接")
                for conn in idle_connections[:close_count]:
                    await self._close_connection(conn)

    async def _repair_unhealthy_connections(self):
        """修复不健康的连接"""
        async with self._lock:
            for conn in self._pool:
                if not await self.validate_connection(conn):
                    conn.health_check_failures += 1

                    # 如果连接健康检查连续失败，则关闭连接
                    if conn.health_check_failures >= self.config.connection_retry:
                        self.config.logger.warning(f"关闭不健康的连接")
                        await self._close_connection(conn)

    async def _ensure_min_connections(self):
        """确保连接池中至少有最小数量的连接"""
        async with self._lock:
            current_count = len(self._pool)
            if current_count < self.config.min_connections:
                needed = self.config.min_connections - current_count
                self.config.logger.info(f"创建 {needed} 个新连接以维持最小连接数")

                for _ in range(needed):
                    try:
                        new_conn = await self.create_connection()
                        if new_conn:
                            self._pool.append(new_conn)
                    except Exception as e:
                        self.config.logger.error(f"创建新连接失败: {e}")

    async def close(self):
        """
        优雅地关闭连接池及所有连接
        """
        if self._closing:
            return

        self._closing = True
        self.config.logger.info("正在关闭连接池...")

        # 取消维护任务
        if self._maintenance_task:
            self._maintenance_task.cancel()
            try:
                await self._maintenance_task
            except asyncio.CancelledError:
                pass

        # 关闭所有连接
        async with self._lock:
            for conn in list(self._pool):
                await self._close_connection(conn)

        self._initialized = False
        self.config.logger.info("连接池已关闭")

    async def check_health(self) -> Dict[str, Any]:
        """
        检查连接池健康状态

        Returns:
            包含连接池健康信息的字典
        """
        async with self._lock:
            total = len(self._pool)
            in_use = sum(1 for conn in self._pool if conn.in_use)
            idle = total - in_use
            healthy_conns = sum(
                1 for conn in self._pool if await self.validate_connection(conn)
            )

            return {
                "total_connections": total,
                "in_use_connections": in_use,
                "idle_connections": idle,
                "healthy_connections": healthy_conns,
                "max_connections": self.config.max_connections,
                "min_connections": self.config.min_connections,
                "status": (
                    "healthy"
                    if self._initialized and not self._closing and healthy_conns > 0
                    else "unhealthy"
                ),
            }


__all__ = ["PoolConfig", "BasePool", "Entity"]
